[mlir][affine] fix the issue of celidiv mul ceildiv expression not satisfying commutative #109382

lipracer · 2024-09-20T07:03:49Z

Fixes #107508

llvmbot · 2024-09-20T07:04:24Z

@llvm/pr-subscribers-llvm-adt
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-affine

Author: long.chen (lipracer)

Changes

Fixs #107508

Full diff: https://github.com/llvm/llvm-project/pull/109382.diff

3 Files Affected:

(modified) llvm/include/llvm/ADT/ScopeExit.h (+6-5)
(modified) mlir/lib/IR/AffineExpr.cpp (+115-49)
(modified) mlir/test/Dialect/Affine/simplify-structures.mlir (+19-3)

diff --git a/llvm/include/llvm/ADT/ScopeExit.h b/llvm/include/llvm/ADT/ScopeExit.h
index 2f13fb65d34d80..7e126479df3a14 100644
--- a/llvm/include/llvm/ADT/ScopeExit.h
+++ b/llvm/include/llvm/ADT/ScopeExit.h
@@ -31,13 +31,14 @@ template <typename Callable> class scope_exit {
   template <typename Fp>
   explicit scope_exit(Fp &&F) : ExitFunction(std::forward<Fp>(F)) {}
 
-  scope_exit(scope_exit &&Rhs)
-      : ExitFunction(std::move(Rhs.ExitFunction)), Engaged(Rhs.Engaged) {
-    Rhs.release();
-  }
+  scope_exit(scope_exit &&Rhs) { *this = std::move(Rhs); }
   scope_exit(const scope_exit &) = delete;
-  scope_exit &operator=(scope_exit &&) = delete;
   scope_exit &operator=(const scope_exit &) = delete;
+  scope_exit &operator=(scope_exit &&Rhs) {
+    Engaged = std::exchange(Rhs.Engaged, false);
+    ExitFunction = std::exchange(Rhs.ExitFunction, {});
+    return *this;
+  }
 
   void release() { Engaged = false; }
 
diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index fc7ede279643ed..84af1f11045d6d 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -9,6 +9,8 @@
 #include <cmath>
 #include <cstdint>
 #include <limits>
+#include <numeric>
+#include <optional>
 #include <utility>
 
 #include "AffineExprDetail.h"
@@ -18,9 +20,8 @@
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/MathExtras.h"
-#include <numeric>
-#include <optional>
 
 using namespace mlir;
 using namespace mlir::detail;
@@ -362,54 +363,119 @@ static bool isDivisibleBySymbol(AffineExpr expr, unsigned symbolPos,
   assert((opKind == AffineExprKind::Mod || opKind == AffineExprKind::FloorDiv ||
           opKind == AffineExprKind::CeilDiv) &&
          "unexpected opKind");
-  switch (expr.getKind()) {
-  case AffineExprKind::Constant:
-    return cast<AffineConstantExpr>(expr).getValue() == 0;
-  case AffineExprKind::DimId:
-    return false;
-  case AffineExprKind::SymbolId:
-    return (cast<AffineSymbolExpr>(expr).getPosition() == symbolPos);
-  // Checks divisibility by the given symbol for both operands.
-  case AffineExprKind::Add: {
-    AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
-    return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind) &&
-           isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, opKind);
-  }
-  // Checks divisibility by the given symbol for both operands. Consider the
-  // expression `(((s1*s0) floordiv w) mod ((s1 * s2) floordiv p)) floordiv s1`,
-  // this is a division by s1 and both the operands of modulo are divisible by
-  // s1 but it is not divisible by s1 always. The third argument is
-  // `AffineExprKind::Mod` for this reason.
-  case AffineExprKind::Mod: {
-    AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
-    return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos,
-                               AffineExprKind::Mod) &&
-           isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos,
-                               AffineExprKind::Mod);
-  }
-  // Checks if any of the operand divisible by the given symbol.
-  case AffineExprKind::Mul: {
-    AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
-    return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, opKind) ||
-           isDivisibleBySymbol(binaryExpr.getRHS(), symbolPos, opKind);
-  }
-  // Floordiv and ceildiv are divisible by the given symbol when the first
-  // operand is divisible, and the affine expression kind of the argument expr
-  // is same as the argument `opKind`. This can be inferred from commutative
-  // property of floordiv and ceildiv operations and are as follow:
-  // (exp1 floordiv exp2) floordiv exp3 = (exp1 floordiv exp3) floordiv exp2
-  // (exp1 ceildiv exp2) ceildiv exp3 = (exp1 ceildiv exp3) ceildiv expr2
-  // It will fail if operations are not same. For example:
-  // (exps1 ceildiv exp2) floordiv exp3 can not be simplified.
-  case AffineExprKind::FloorDiv:
-  case AffineExprKind::CeilDiv: {
-    AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
-    if (opKind != expr.getKind())
-      return false;
-    return isDivisibleBySymbol(binaryExpr.getLHS(), symbolPos, expr.getKind());
-  }
+  std::vector<std::tuple<AffineExpr, unsigned, AffineExprKind,
+                         llvm::detail::scope_exit<std::function<void(void)>>>>
+      stack;
+  stack.emplace_back(expr, symbolPos, opKind, []() {});
+  bool result = false;
+
+  while (!stack.empty()) {
+    AffineExpr expr = std::get<0>(stack.back());
+    unsigned symbolPos = std::get<1>(stack.back());
+    AffineExprKind opKind = std::get<2>(stack.back());
+
+    switch (expr.getKind()) {
+    case AffineExprKind::Constant: {
+      // Note: Assignment must occur before pop, which will affect whether it
+      // enters other execution branches.
+      result = cast<AffineConstantExpr>(expr).getValue() == 0;
+      stack.pop_back();
+      break;
+    }
+    case AffineExprKind::DimId: {
+      result = false;
+      stack.pop_back();
+      break;
+    }
+    case AffineExprKind::SymbolId: {
+      result = cast<AffineSymbolExpr>(expr).getPosition() == symbolPos;
+      stack.pop_back();
+      break;
+    }
+    // Checks divisibility by the given symbol for both operands.
+    case AffineExprKind::Add: {
+      AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
+      stack.emplace_back(binaryExpr.getLHS(), symbolPos, opKind,
+                         [&stack, &result, binaryExpr, symbolPos, opKind]() {
+                           if (result) {
+                             stack.emplace_back(
+                                 binaryExpr.getRHS(), symbolPos, opKind,
+                                 [&stack]() { stack.pop_back(); });
+                           } else {
+                             stack.pop_back();
+                           }
+                         });
+      break;
+    }
+    // Checks divisibility by the given symbol for both operands. Consider the
+    // expression `(((s1*s0) floordiv w) mod ((s1 * s2) floordiv p)) floordiv
+    // s1`, this is a division by s1 and both the operands of modulo are
+    // divisible by s1 but it is not divisible by s1 always. The third argument
+    // is `AffineExprKind::Mod` for this reason.
+    case AffineExprKind::Mod: {
+      AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
+      stack.emplace_back(binaryExpr.getLHS(), symbolPos, AffineExprKind::Mod,
+                         [&stack, &result, binaryExpr, symbolPos, opKind]() {
+                           if (result) {
+                             stack.emplace_back(
+                                 binaryExpr.getRHS(), symbolPos,
+                                 AffineExprKind::Mod,
+                                 [&stack]() { stack.pop_back(); });
+                           } else {
+                             stack.pop_back();
+                           }
+                         });
+      break;
+    }
+    // Checks if any of the operand divisible by the given symbol.
+    case AffineExprKind::Mul: {
+      AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
+      stack.emplace_back(binaryExpr.getLHS(), symbolPos, opKind,
+                         [&stack, &result, binaryExpr, symbolPos, opKind]() {
+                           if (!result) {
+                             stack.emplace_back(
+                                 binaryExpr.getRHS(), symbolPos, opKind,
+                                 [&stack]() { stack.pop_back(); });
+                           } else {
+                             stack.pop_back();
+                           }
+                         });
+      break;
+    }
+    // Floordiv and ceildiv are divisible by the given symbol when the first
+    // operand is divisible, and the affine expression kind of the argument expr
+    // is same as the argument `opKind`. This can be inferred from commutative
+    // property of floordiv and ceildiv operations and are as follow:
+    // (exp1 floordiv exp2) floordiv exp3 = (exp1 floordiv exp3) floordiv exp2
+    // (exp1 ceildiv exp2) ceildiv exp3 = (exp1 ceildiv exp3) ceildiv expr2
+    // It will fail 1.if operations are not same. For example:
+    // (exps1 ceildiv exp2) floordiv exp3 can not be simplified. 2.if there is a
+    // multiplication operation in the expression. For example:
+    //  (exps1 ceildiv exp2) mul exp3 ceildiv exp4 can not be simplified.
+    case AffineExprKind::FloorDiv:
+    case AffineExprKind::CeilDiv: {
+      AffineBinaryOpExpr binaryExpr = cast<AffineBinaryOpExpr>(expr);
+      if (opKind != expr.getKind()) {
+        result = false;
+        stack.pop_back();
+        break;
+      }
+      if (llvm::any_of(stack, [](auto &it) {
+            return std::get<0>(it).getKind() == AffineExprKind::Mul;
+          })) {
+        result = false;
+        stack.pop_back();
+        break;
+      }
+
+      stack.emplace_back(binaryExpr.getLHS(), symbolPos, expr.getKind(),
+                         [&stack]() { stack.pop_back(); });
+      break;
+    }
+      llvm_unreachable("Unknown AffineExpr");
+    }
   }
-  llvm_unreachable("Unknown AffineExpr");
+  return result;
 }
 
 /// Divides the given expression by the given symbol at position `symbolPos`. It
diff --git a/mlir/test/Dialect/Affine/simplify-structures.mlir b/mlir/test/Dialect/Affine/simplify-structures.mlir
index 92d3d86bc93068..d1f34f20fa5dad 100644
--- a/mlir/test/Dialect/Affine/simplify-structures.mlir
+++ b/mlir/test/Dialect/Affine/simplify-structures.mlir
@@ -308,10 +308,26 @@ func.func @semiaffine_ceildiv(%arg0: index, %arg1: index) -> index {
 }
 
 // Tests the simplification of a semi-affine expression with a nested ceildiv operation and further simplifications after performing ceildiv.
-// CHECK-LABEL: func @semiaffine_composite_floor
-func.func @semiaffine_composite_floor(%arg0: index, %arg1: index) -> index {
+// CHECK-LABEL: func @semiaffine_composite_ceildiv
+func.func @semiaffine_composite_ceildiv(%arg0: index, %arg1: index) -> index {
+  %a = affine.apply affine_map<(d0)[s0] ->((((s0 * 2) ceildiv 4) + s0 * 42) ceildiv s0)> (%arg0)[%arg1]
+  // CHECK:       %[[CST:.*]] = arith.constant 43
+  return %a : index
+}
+
+// Tests the do not simplification of a semi-affine expression with a nested ceildiv-mul-ceildiv operation.
+// CHECK-LABEL: func @semiaffine_composite_ceildiv
+func.func @semiaffine_composite_ceildiv_mul_ceildiv(%arg0: index, %arg1: index) -> index {
   %a = affine.apply affine_map<(d0)[s0] ->(((((s0 * 2) ceildiv 4) * 5) + s0 * 42) ceildiv s0)> (%arg0)[%arg1]
-  // CHECK:       %[[CST:.*]] = arith.constant 47
+  // CHECK-NOT:       arith.constant
+  return %a : index
+}
+
+// Tests the do not simplification of a semi-affine expression with a nested floordiv_mul_floordiv operation
+// CHECK-LABEL: func @semiaffine_composite_floordiv
+func.func @semiaffine_composite_floordiv_mul_floordiv(%arg0: index, %arg1: index) -> index {
+  %a = affine.apply affine_map<(d0)[s0] ->(((((s0 * 2) floordiv 4) * 5) + s0 * 42) floordiv s0)> (%arg0)[%arg1]
+  // CHECK-NOT:       arith.constant
   return %a : index
 }

jreiffers · 2024-09-20T11:09:40Z

Hey, I see you're still actively changing this. Let me know when it's ready for review.

lipracer · 2024-09-20T11:14:06Z

Hey, I see you're still actively changing this. Let me know when it's ready for review.

Ready for review now.

github-actions · 2024-09-28T12:17:01Z

✅ With the latest revision this PR passed the C/C++ code formatter.

…ot satisfying commutative Fixes llvm#107508

mlir/lib/IR/AffineExpr.cpp

jreiffers · 2024-10-04T06:19:57Z

mlir/lib/IR/AffineExpr.cpp

  return static_cast<ImplType *>(expr)->position;
 }

+/// A manually managed stack used to convert recursive function calls into


Why not just use AffineExpr::walk? It'll do a post-order traversal, so you could just have a stack of partial results. Unless I'm still missing something :-).

This requires pre order traversal. Perhaps we can refactor Affine Visitor to support pre order traversal.

Can you give me a small example/explanation where pre-order traversal is needed?

Sorry,I didn't express clearly. The implementation of this isDivisibleBySymbol requires a pre-order traversal.
eg:

case add: visit(lhs) and vist(rhs) case mul: visit(lhs) or vist(rhs)

base on the current expr type, using a control traversal approach, perhaps returning the interrupting and boolean types in the visitor can meet this requirement, but it also needs to be reimplemented.

In a post-order traversal with a stack for intermediates, you'd have the results for lhs and rhs on the top of the stack when you visit an add, right? Can you give me a more complete example where this doesn't work?

I see no problem with using an explicit stack in addition to the implicit one in ::walk.

If you're concerned about stack overflows, the correct thing to do is to refactor ::walk to not be recursive, since that function is used in many places, so fixing this particular location isn't really sufficient if stack overflows are an issue with ::walk.

Refactoring MLIR walk() to not be recursive is a long-lasting issue: patch welcome to fix this!

Ok.I will switch to a ‘walk’ approach and then refactor AffineExprVisitor in another patch using a stack to avoid recursive calls.

Why don't use walk,
This is walk implementation.At this point,lhs and rhs have already been traversed.and I don't know where to do the push stack.

@joker-eph The walk here is the traversal method of AffineExpr. When you mention mlir:: walk, do you mean about mlir::Operation?

llvmbot added mlir:core MLIR Core Infrastructure mlir:affine mlir llvm:adt labels Sep 20, 2024

lipracer changed the title ~~[mlir][affine] fix the issue of celidiv mul childiv expression not sa…~~ [mlir][affine] fix the issue of celidiv mul childiv expression not satisfying commutative Sep 20, 2024

lipracer requested review from jreiffers and kazutakahirata September 20, 2024 07:04

lipracer force-pushed the fix-107508 branch 4 times, most recently from af8be8c to 21587f8 Compare September 20, 2024 10:42

lipracer force-pushed the fix-107508 branch from 0ed38c8 to 4410645 Compare September 20, 2024 14:41

lipracer changed the title ~~[mlir][affine] fix the issue of celidiv mul childiv expression not satisfying commutative~~ [mlir][affine] fix the issue of celidiv mul ceildiv expression not satisfying commutative Sep 22, 2024

lipracer added the awaiting-review Has pending Phabricator review label Sep 23, 2024

lipracer force-pushed the fix-107508 branch from 4410645 to 389dc01 Compare September 24, 2024 08:57

MaheshRavishankar requested review from Groverkss and matthias-springer September 24, 2024 13:43

lipracer force-pushed the fix-107508 branch from 389dc01 to e096a49 Compare September 28, 2024 12:13

lipracer force-pushed the fix-107508 branch from e096a49 to 437549c Compare September 28, 2024 12:22

lipracer added 4 commits September 28, 2024 11:11

[mlir][affine] fix the issue of ceildiv-mul-ceildiv form expression n…

1858636

…ot satisfying commutative Fixes llvm#107508

refine

c758224

refine

dee51cd

fix ci fail on windows platform

c70e794

lipracer force-pushed the fix-107508 branch from 437549c to c70e794 Compare September 28, 2024 15:13

jreiffers reviewed Sep 30, 2024

View reviewed changes

mlir/lib/IR/AffineExpr.cpp Outdated Show resolved Hide resolved

refine

514d78c

lipracer requested review from jreiffers October 4, 2024 04:18

jreiffers reviewed Oct 4, 2024

View reviewed changes

lipracer closed this Oct 5, 2024

[mlir][affine] fix the issue of celidiv mul ceildiv expression not satisfying commutative #109382

[mlir][affine] fix the issue of celidiv mul ceildiv expression not satisfying commutative #109382

Uh oh!

Conversation

lipracer commented Sep 20, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Sep 20, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jreiffers commented Sep 20, 2024

Uh oh!

lipracer commented Sep 20, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Sep 28, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

lipracer commented Sep 20, 2024 •

edited

Loading

llvmbot commented Sep 20, 2024 •

edited

Loading

lipracer commented Sep 20, 2024 •

edited

Loading

github-actions bot commented Sep 28, 2024 •

edited

Loading